
#include <map>
#include <sstream>
#include <iostream>
#include <filesystem>
#include <vector>
#include <fstream>
#include <string>
#include <argp.h>
#include <unistd.h>
#include <time.h>
#include <dirent.h>
#include <elfutils/libdwfl.h>
#include <linux/prctl.h>
#include <linux/seccomp.h>
#include <linux/capability.h>
#include <linux/securebits.h>
#include <asm/unistd_64.h>
#include <bcc/bcc_syms.h>

#include "cap.bpf.h"
#include "cap.bpf.skel.h"

static std::string usage = "Usage: [-p pid] [-u uid] [-v verbose] [-d debug libbpf print]";

static bool g_debug = false;
static bool g_verbose = false;

/** ref: https://en.wikipedia.org/wiki/ANSI_escape_code */
#define ANSI_RED     "\x1b[31m"
#define ANSI_GREEN   "\x1b[32m"
#define ANSI_YELLOW  "\x1b[33m"
#define ANSI_BLUE    "\x1b[34m"
#define ANSI_WHITE   "\x1b[37m"
#define ANSI_RESET   "\x1b[0m"

#define printf_syscall(fmt, ...)                                printf("syscall: " fmt "\n", ##__VA_ARGS__)
#define printf_syscall_args(name, fmt, ...)                     printf(name ": " fmt "\n", ##__VA_ARGS__)
#define printf_syscall_args_positive(name, fmt, positive)       if (positive > 0) printf_syscall_args(name, fmt, positive)

/**
 * 缓存进程信息不可以用bcc处理
 * bcc的符号缓存必须和pid绑定 不可以fork时继承父进程的符号缓存
*/
class BCCSymcache {
public:
    BCCSymcache(pid_t pid)
        : pid(pid) {
        struct bcc_symbol_option symopts;

        memset(&symopts, 0, sizeof(symopts));
        symopts.use_debug_file = 1;
        symopts.check_debug_file_crc = 1;
        symopts.use_symbol_type = BCC_SYM_ALL_TYPES;

        psyms = bcc_symcache_new(pid, &symopts);
    }

    ~BCCSymcache() {
        bcc_free_symcache(psyms, pid);
    }

    void refresh() {
        bcc_symcache_refresh(psyms);
    }

    void* get() { return psyms; }

private:
    pid_t pid;
    void* psyms;
};

class BCCUstackResolver {
public:
    BCCUstackResolver() {
        DIR* dir = opendir("/proc");
        if (!dir) {
            std::cerr << "failed to open /proc" << std::endl;
            return;
        }

        while (dirent* d = readdir(dir)) {
            int pid = atoi(d->d_name);
            if (pid > 0 || (pid == 0 && d->d_name[0] == '0')) {
                init_process(pid);
            }
        }

        closedir(dir);
    }

    ~BCCUstackResolver() {
        exe_sym_.clear();
    }

    std::string resolve_process(std::vector<uint64_t> addrs, int pid, bool show_offset = true, bool show_module = true) {
        struct bcc_symbol usym;
        std::ostringstream symbol;

        if (exe_sym_.find(pid) == exe_sym_.end())
            init_process(pid);

        for (auto addr : addrs) {
            if (exe_sym_[pid]->get() && bcc_symcache_resolve(exe_sym_[pid]->get(), addr, &usym) == 0) {
                // symbol << usym.demangle_name;
                symbol << usym.name;

                if (show_module)
                    symbol << "+" << usym.offset << " (" << usym.module << ")";
            }
            else {
                symbol << "unknown";
            }

            symbol << " " << (void *)addr << std::endl;
        }

        return symbol.str();
    }

    /**
     * 不判断pid是否已缓存 直接init
    */
    void init_process(int pid) {
        exe_sym_[pid] = std::make_unique<BCCSymcache>(pid);
    }

    void exit_process(int pid) {
        if (exe_sym_.find(pid) != exe_sym_.end()) {
            exe_sym_.erase(pid);
        }
    }

    void refresh_process(int pid) {
        if (exe_sym_.find(pid) != exe_sym_.end())
            exe_sym_[pid]->refresh();
        else
            init_process(pid);
    }

    void fork_process(int pid, int ppid) {
        exit_process(pid);
        init_process(pid);
    }

private:
    std::map<int, std::unique_ptr<BCCSymcache>> exe_sym_; // exe -> (pid, cache)
};

static std::vector<std::string> secbits2str = {
    "SECBIT_NOROOT",
    "SECBIT_NOROOT_LOCKED",
    "SECBIT_NO_SETUID_FIXUP",
    "SECBIT_NO_SETUID_FIXUP_LOCKED",
    "SECBIT_KEEP_CAPS",
    "SECBIT_KEEP_CAPS_LOCKED",
    "SECBIT_NO_CAP_AMBIENT_RAISE",
    "SECBIT_NO_CAP_AMBIENT_RAISE_LOCKED",
};

static std::vector<std::string> cap2str = {
    "CAP_CHOWN",
    "CAP_DAC_OVERRIDE",
    "CAP_DAC_READ_SEARCH",
    "CAP_FOWNER",
    "CAP_FSETID",
    "CAP_KILL",
    "CAP_SETGID",
    "CAP_SETUID",
    "CAP_SETPCAP",
    "CAP_LINUX_IMMUTABLE",
    "CAP_NET_BIND_SERVICE",
    "CAP_NET_BROADCAST",
    "CAP_NET_ADMIN",
    "CAP_NET_RAW",
    "CAP_IPC_LOCK",
    "CAP_IPC_OWNER",
    "CAP_SYS_MODULE",
    "CAP_SYS_RAWIO",
    "CAP_SYS_CHROOT",
    "CAP_SYS_CHROOT",
    "CAP_SYS_PTRACE",
    "CAP_SYS_PACCT",
    "CAP_SYS_ADMIN",
    "CAP_SYS_BOOT",
    "CAP_SYS_NICE",
    "CAP_SYS_RESOURCE",
    "CAP_SYS_TIME",
    "CAP_SYS_TTY_CONFIG",
    "CAP_MKNOD",
    "CAP_LEASE",
    "CAP_AUDIT_CONTROL",
    "CAP_AUDIT_WRITE",
    "CAP_SETFCAP",
    "CAP_MAC_OVERRIDE",
    "CAP_MAC_ADMIN",
    "CAP_SYSLOG",
    "CAP_WAKE_ALARM",
    "CAP_BLOCK_SUSPEND",
    "CAP_AUDIT_READ",
    "CAP_PERFMON",
    "CAP_BPF",
    "CAP_CHECKPOINT_RESTORE",
};

void nsec_to_hms(uint64_t nsec, struct tm *result, long *nsec_part) {
    // 将纳秒转换为秒和纳秒部分
    time_t seconds = nsec / 1000000000;
    *nsec_part = nsec % 1000000000;

    // 使用gmtime_r将时间转换为UTC时间的时、分、秒结构
    gmtime_r(&seconds, result);
}

static void print_event(struct s_event* e) {
    switch (e->nr) {
        case __NR_clone:
            printf_syscall("clone");
            break;

        case __NR_fork:
            printf_syscall("fork");
            break;

        case __NR_execve:
            printf_syscall("execve");
            printf_syscall_args("filename", "%s", e->data.execve.filename);
            break;

        case __NR_execveat:
            printf_syscall("execveat");
            printf_syscall_args("filename", "%s", e->data.execve.filename);
            break;

        case __NR_exit:
            printf_syscall("exit");
            break;

        case __NR_setuid:
            printf_syscall("setuid");
            printf_syscall_args("uid", "%d", e->data.setuid.uid);
            break;

        case __NR_setgid:
            printf_syscall("setgid");
            printf_syscall_args("gid", "%d", e->data.setgid.gid);
            break;

        case __NR_setpgid:
            printf_syscall("setpgid");
            printf_syscall_args("pid", "%d", e->data.setgid.pid);
            printf_syscall_args("pgid", "%d", e->data.setgid.pgid);
            break;

        case __NR_setreuid:
            printf_syscall("setreuid");
            printf_syscall_args("uid", "%d", e->data.setuid.uid);
            printf_syscall_args("euid", "%d", e->data.setuid.euid);
            break;

        case __NR_setregid:
            printf_syscall("setregid");
            printf_syscall_args("rgid", "%d", e->data.setgid.rgid);
            printf_syscall_args("egid", "%d", e->data.setgid.egid);
            break;

        case __NR_setresuid:
            printf_syscall("setresuid");
            printf_syscall_args("uid", "%d", e->data.setuid.uid);
            printf_syscall_args("euid", "%d", e->data.setuid.euid);
            printf_syscall_args("suid", "%d", e->data.setuid.suid);
            break;

        case __NR_setresgid:
            printf_syscall("setresgid");
            printf_syscall_args("rgid", "%d", e->data.setgid.rgid);
            printf_syscall_args("egid", "%d", e->data.setgid.egid);
            printf_syscall_args("sgid", "%d", e->data.setgid.sgid);
            break;

        case __NR_setfsuid:
            printf_syscall("setfsuid");
            printf_syscall_args("uid", "%d", e->data.setuid.uid);
            break;

        case __NR_setfsgid:
            printf_syscall("setfsgid");
            printf_syscall_args("gid", "%d", e->data.setgid.gid);
            break;

        case __NR_capset:
            printf_syscall("capset");

            if (e->data.capset.hdrp.version == _LINUX_CAPABILITY_VERSION_1)
                printf_syscall_args("hdrp.version", "%s", "LINUX_CAPABILITY_VERSION_1");
            else if (e->data.capset.hdrp.version == _LINUX_CAPABILITY_VERSION_2)
                printf_syscall_args("hdrp.version", "%s", "LINUX_CAPABILITY_VERSION_2");
            else if (e->data.capset.hdrp.version == _LINUX_CAPABILITY_VERSION_3)
                printf_syscall_args("hdrp.version", "%s", "LINUX_CAPABILITY_VERSION_3");
            else
                printf_syscall_args("hdrp.version", "%x", e->data.capset.hdrp.version);

            printf_syscall_args("hdrp.pid", "%d", e->data.capset.hdrp.pid);
            printf_syscall_args("effective", "0x%x", e->data.capset.datap.effective);
            printf_syscall_args("permitted", "0x%x", e->data.capset.datap.permitted);
            printf_syscall_args("inheritable", "0x%x", e->data.capset.datap.inheritable);
            break;

        case __NR_prctl:
            printf_syscall("prctl");

            if (e->data.prctl.option == PR_SET_KEEPCAPS) {
                printf_syscall_args("option", "PR_SET_KEEPCAPS");
                printf_syscall_args("arg2", "%d", e->data.prctl.arg2);
            }
            else if (e->data.prctl.option == PR_SET_SECCOMP) {
                printf_syscall_args("option", "PR_SET_SECCOMP");
                if (e->data.prctl.arg2 == SECCOMP_MODE_DISABLED)
                    printf_syscall_args("arg2", "SECCOMP_MODE_DISABLED");
                else if (e->data.prctl.arg2 == SECCOMP_MODE_STRICT)
                    printf_syscall_args("arg2", "SECCOMP_MODE_STRICT");
                else if (e->data.prctl.arg2 == SECCOMP_MODE_FILTER)
                    printf_syscall_args("arg2", "SECCOMP_MODE_FILTER");
                else {
                    printf_syscall_args_positive("arg2", "%d", e->data.prctl.arg2);
                }

                printf_syscall_args_positive("arg3", "%d", e->data.prctl.arg3);
                printf_syscall_args_positive("arg4", "%d", e->data.prctl.arg4);
                printf_syscall_args_positive("arg5", "%d", e->data.prctl.arg5);

            }
            else if (e->data.prctl.option == PR_SET_SECUREBITS) {
                printf_syscall_args("option", "PR_SET_SECUREBITS");
                printf_syscall_args("arg2", "%x", e->data.prctl.arg2);
            }
            else if (e->data.prctl.option == PR_CAP_AMBIENT) {
                printf_syscall_args("option", "PR_CAP_AMBIENT");

                if (e->data.prctl.arg2 == PR_CAP_AMBIENT_IS_SET)
                    printf_syscall_args("arg2", "PR_CAP_AMBIENT_IS_SET");
                else if (e->data.prctl.arg2 == PR_CAP_AMBIENT_RAISE)
                    printf_syscall_args("arg2", "PR_CAP_AMBIENT_RAISE");
                else if (e->data.prctl.arg2 == PR_CAP_AMBIENT_LOWER)
                    printf_syscall_args("arg2", "PR_CAP_AMBIENT_LOWER");
                else if (e->data.prctl.arg2 == PR_CAP_AMBIENT_CLEAR_ALL)
                    printf_syscall_args("arg2", "PR_CAP_AMBIENT_CLEAR_ALL");
                else {
                    printf_syscall_args_positive("arg2", "%d", e->data.prctl.arg2);
                }

                if (e->data.prctl.arg3 < (int)cap2str.size())
                    printf_syscall_args("arg3", "%s", cap2str[e->data.prctl.arg3].c_str());
                else {
                    printf_syscall_args_positive("arg3", "%d", e->data.prctl.arg3);
                }

                printf_syscall_args_positive("arg4", "%d", e->data.prctl.arg4);
                printf_syscall_args_positive("arg5", "%d", e->data.prctl.arg5);
            }
            else {
                printf_syscall_args_positive("option", "%d", e->data.prctl.arg2);
                printf_syscall_args_positive("arg2", "%d", e->data.prctl.arg2);
                printf_syscall_args_positive("arg3", "%d", e->data.prctl.arg3);
                printf_syscall_args_positive("arg4", "%d", e->data.prctl.arg4);
                printf_syscall_args_positive("arg5", "%d", e->data.prctl.arg5);
            }
            break;

        case __NR_clone3:
            printf_syscall("clone3");
            break;

        default:
            printf_syscall("unknown: %d", e->nr);
            break;
    }

    printf_syscall_args("retval", "%d", e->retvel);
}

static void print_cap(struct s_cap* cap) {
    printf_syscall_args("securebits", "0x%08x", cap->securebits);
    printf_syscall_args("cap_inheritable", "0x%016llx", cap->cap_inheritable);
    printf_syscall_args("cap_permitted", "0x%016llx", cap->cap_permitted);
    printf_syscall_args("cap_effective", "0x%016llx", cap->cap_effective);
    printf_syscall_args("cap_bset", "0x%016llx", cap->cap_bset);
    printf_syscall_args("cap_ambient", "0x%016llx", cap->cap_ambient);
}

static void print_cap_verbose_diff(__u64 before, __u64 after, std::vector<std::string> diffstr) {
    __u64 added = 0;
    __u64 removed = 0;
    int cnt = 0;

    for (std::size_t i = 0; i < diffstr.size(); i++) {
        if ((before & ((__u64)0x1 << i)) != (after & ((__u64)0x1 << i))) {
            if (before & ((__u64)0x1 << i))
                removed |= ((__u64)0x1 << i);
            else
                added |= ((__u64)0x1 << i);

            cnt += 1;
        }
    }

    {
        int j = 0;
        for (std::size_t i = 0; i < diffstr.size(); i++) {
            if (added & ((__u64)0x1 << i)) {
                printf(ANSI_GREEN "%s" ANSI_RESET, diffstr[i].c_str());
                j++;

                if (j != (cnt - 1))
                printf(", ");
            }
            else if (removed & ((__u64)0x1 << i)) {
                printf(ANSI_RED "%s" ANSI_RESET, diffstr[i].c_str());
                j++;

                if (j != (cnt - 1))
                printf(", ");
            }

            if (j == (cnt - 1))
                break;
        }
    }
}

static void print_hex_diff(__u64 before, __u64 after, __u8 size) {
    printf("0x");

    for (int i = (size * 2 - 1); i >= 0; i--) {
        __u64 target = (__u64)0xf << (4 * i);
        if ((target & before) == (target & after))
            printf("%x", (__u8)((target & before) >> (4 * i)));
        else
            printf(ANSI_RED "%x" ANSI_RESET, (__u8)((target & after) >> (4 * i)));
    }
}

static void print_cap_item_diff(std::string name, __u64 before, __u64 after, __u8 size) {
    if (before == after)
        printf("%s: ", name.c_str());
    else
        printf(ANSI_RED "%s: " ANSI_RESET, name.c_str());
    
    print_hex_diff(before, after, size);

    if (before == after || !g_verbose) {         // 只打印变化的数字
        printf("\n");
        return;
    }

    printf(", (");

    if (size == 4)
        print_cap_verbose_diff(before, after, secbits2str);
    else
        print_cap_verbose_diff(before, after, cap2str);

    printf(")\n");
    return;
}

static void print_cap_diff(struct s_cap* cap_before, struct s_cap* cap_after) {
    print_cap_item_diff("securebits", cap_before->securebits, cap_after->securebits, sizeof(cap_before->securebits));
    print_cap_item_diff("cap_inheritable", cap_before->cap_inheritable, cap_after->cap_inheritable, sizeof(cap_before->cap_inheritable));
    print_cap_item_diff("cap_permitted", cap_before->cap_permitted, cap_after->cap_permitted, sizeof(cap_before->cap_permitted));
    print_cap_item_diff("cap_effective", cap_before->cap_effective, cap_after->cap_effective, sizeof(cap_before->cap_effective));
    print_cap_item_diff("cap_bset", cap_before->cap_bset, cap_after->cap_bset, sizeof(cap_before->cap_bset));
    print_cap_item_diff("cap_ambient", cap_before->cap_ambient, cap_after->cap_ambient, sizeof(cap_before->cap_ambient));
}

static int libbpf_print(enum libbpf_print_level level, const char *format, va_list args)
{
    if (g_debug)
        return vfprintf(stderr, format, args);
    else
        return 0;
}

static void perf_buffer_sample(void *ctx, int cpu, void *data, __u32 size) {
    static BCCUstackResolver bcc_resolver;

    if (size == (sizeof(struct s_maps_event) + 4)) {
        struct s_maps_event* event = (struct s_maps_event*)data;

        switch (event->nr)
        {
        case -1:
            bcc_resolver.exit_process(event->pid);
            break;

        default:
            std::cerr << "unknown maps event nr = " << event->nr << std::endl;
            break;
        }
        return;
    }

    struct s_event* event = (struct s_event*)data;

    /**** 权能改变 输出 *****/
    struct tm timeinfo;
    long nsec_part;

    nsec_to_hms(event->nsec, &timeinfo, &nsec_part);

    // 打印时、分、秒和纳秒部分
    printf(ANSI_GREEN"Time: %02d:%02d:%02d.%09ld",
        timeinfo.tm_hour,
        timeinfo.tm_min,
        timeinfo.tm_sec,
        nsec_part);

    printf(", uid = %d, pid: %d, tgid: %d, ppid: %d, comm: %s\n", event->uid, event->pid, event->tgid, event->ppid, event->comm);
    print_event(event);
    printf(ANSI_RESET);
    print_cap(&event->cap_before);
    printf("----------------------------------\n");
    print_cap_diff(&event->cap_before, &event->cap_after);

    if (event->ustack_sz > 0) {
        printf("----------------------------------\n");
        std::cout << bcc_resolver.resolve_process(std::vector<uint64_t>(event->ustack, event->ustack + (event->ustack_sz / 8)),
                                            event->pid, true, true)
                << std::endl;
    }

    printf("\n");
}

static void perf_buffer_lost(void *ctx, int cpu, __u64 cnt) {
    printf("lost %llu events at cpu %d\n", cnt, cpu);
}

static void parse_args(int args, char* argv[], struct s_filter* filter) {
    int opt;

    while ((opt = getopt(args, argv, "dvp:u:h")) != -1) {
        switch(opt) {
            case 'p':
                filter->pid = atoi(optarg);
                break;
            case 'u':
                filter->uid = atoi(optarg);
                break;
            
            case 'd':
                g_debug = true;
                break;
            
            case 'v':
                g_verbose = true;
                break;

            default:
                std::cerr << usage << std::endl;
                exit(-1);
        }
    }

    if (filter->pid >= 0)
        printf("filter pid: %d\n", filter->pid);
    if (filter->uid >= 0)
        printf("filter uid: %lld\n", filter->uid);
}

int main(int args, char* argv[]) {
    int err;

    libbpf_set_print(libbpf_print);

    struct cap_bpf* capsnoop = cap_bpf__open_and_load();
    if (!capsnoop) {
        fprintf(stderr, "failed to open and load bpf object\n");
        return -1;
    }

    /** 设置过滤条件 */
    parse_args(args, argv, &capsnoop->data->filter);

    /** bpf run */
    err = cap_bpf__attach(capsnoop);
    if (err) {
        fprintf(stderr, "failed to attach bpf object\n");
        return -1;
    }

#ifdef ___perf_buffer_new6
    struct perf_buffer* pb = perf_buffer__new(bpf_map__fd(capsnoop->maps.events),
                                              128, perf_buffer_sample, perf_buffer_lost, NULL, NULL);
#else
    struct perf_buffer_opts opts = {
    	.lost_cb = perf_buffer_lost,
	    .sample_cb = perf_buffer_sample,
    };

    struct perf_buffer* pb = perf_buffer__new(bpf_map__fd(capsnoop->maps.events),
                                              128, &opts);
#endif

    err = libbpf_get_error(pb);
	if (err) {
		pb = NULL;
		fprintf(stderr, "failed to open perf buffer: %d\n", err);
		goto cleanup;
	}

	/* Main polling loop */
	while ((err = perf_buffer__poll(pb, 100)) >= 0);
	printf("Error polling perf buffer: %d\n", err);

cleanup:
    perf_buffer__free(pb);
    cap_bpf__destroy(capsnoop);
    return 0;
}
